# Author Zenos
# Create 2025/3/11 下午3:40
import torch
import torch.nn as nn
import torchvision.models as models
import numpy as np


def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size,
                       kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight)


class FiLM(nn.Module):
    def __init__(self, num_classes, feature_dim):
        """
        FiLM 层：根据类别信息调节 CNN 的中间特征

        :param num_classes: 总类别数
        :param feature_dim: 需要调节的特征图通道数
        """
        super(FiLM, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(num_classes, 64),  # 全连接层1
            nn.ReLU(),  # 激活函数
            nn.Linear(64, 2 * feature_dim)  # 全连接层2（输出γ和β）
        )

    def forward(self, features, presence_vector):
        params = self.fc(presence_vector)  # 生成参数
        gamma, beta = torch.chunk(params, 2, dim=1)
        gamma = gamma.unsqueeze(-1).unsqueeze(-1)  # 维度扩展
        beta = beta.unsqueeze(-1).unsqueeze(-1)
        return features * gamma + beta  # 特征调整


class ASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ASPP, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 1, padding=0)
        self.conv2 = nn.Conv2d(in_channels, out_channels, 3, padding=6, dilation=6)
        self.conv3 = nn.Conv2d(in_channels, out_channels, 3, padding=12, dilation=12)
        self.conv4 = nn.Conv2d(in_channels, out_channels, 3, padding=18, dilation=18)
        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
        self.global_pool = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
        self.final_conv = nn.Conv2d(out_channels * 5, out_channels, 1)

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x3 = self.conv3(x)
        x4 = self.conv4(x)
        x5 = self.global_avg_pool(x)
        x5 = self.global_pool(x5)
        x5 = nn.functional.interpolate(x5, size=x4.shape[2:], mode='bilinear', align_corners=True)
        x = torch.cat([x1, x2, x3, x4, x5], dim=1)
        return self.final_conv(x)


class Backbone(nn.Module):
    def __init__(self, num_classes=35):
        super(Backbone, self).__init__()
        resnet = models.resnet34(weights=None)
        self.stage = list(resnet.children())[:4]  # 第一段
        # 输入 [1,320,320]
        # 修改第一层的通道数  输出[1, 64, 160, 160]
        self.stage[0] = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        # 将第四层最大池化改成空洞卷积   输出[1, 64, 80, 80]
        self.stage[3] = nn.Conv2d(64, 64, 2, stride=2, padding=1, bias=False, dilation=2)
        # 加入一层批标准化
        batch = nn.BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.stage.append(batch)
        # # 转换成Sequential
        self.stage = nn.Sequential(*self.stage)

        # FiLM 模块
        self.film_layer = FiLM(num_classes, feature_dim=64)

        # 添加 ResNet 的 4~7 层
        # 加入
        self.layer1 = list(resnet.children())[4]  # layer1   输出[64,80,80]
        self.layer2 = list(resnet.children())[5]  # layer2   输出[64,40,40]
        self.layer3 = list(resnet.children())[6]  # layer3   输出[64,20,20]
        self.layer4 = list(resnet.children())[7]  # layer4   输出[64,10,10]
        self.aspp = ASPP(512, 256)

    def forward(self, x, valid_classes):
        x = self.stage(x)
        x = self.film_layer(x, valid_classes)  # 在 CNN 特征图上应用 FiLM
        x = self.layer1(x)  # [1, 64, 80, 80]
        s1 = x
        x = self.layer2(x)  # [1, 128, 40, 40]
        s2 = x
        x = self.layer3(x)  # [1, 256, 20, 20]
        s3 = x
        x = self.layer4(x)  # [1, 512, 10, 10]
        x = self.aspp(x)
        return x, s1, s2, s3


class UStroke(nn.Module):
    def __init__(self, num_classes=35):
        super(UStroke, self).__init__()
        self.backbone = Backbone(num_classes)
        # 改变通道数
        self.change_aspp = nn.Conv2d(256, num_classes, 1)
        self.scores1 = nn.Conv2d(256, num_classes, 1)
        self.scores2 = nn.Conv2d(128, num_classes, 1)
        self.scores3 = nn.Conv2d(64, num_classes, 1)

        # 2倍上采样
        self.upsample_2x = nn.ConvTranspose2d(
            num_classes, num_classes, 4, 2, 1, bias=False)
        self.upsample_2x.weight.data = bilinear_kernel(
            num_classes, num_classes, 4)

    def forward(self, x, c):
        x, x_64, x_128, x_256 = self.backbone(x,c)
        x_64 = self.scores3(x_64)
        x_128 = self.scores2(x_128)
        x_256 = self.scores1(x_256)
        x = self.change_aspp(x)
        x = self.upsample_2x(x)
        x = x + x_256
        x = self.upsample_2x(x)
        x = x + x_128
        x = self.upsample_2x(x)
        x = x + x_64
        x = self.upsample_2x(x)
        x = self.upsample_2x(x)
        return x


# 测试
if __name__ == '__main__':
    model = UStroke()
    print(model)
    X = torch.randn(1, 1, 320, 320)
    output = model(X)
    print(output.shape)
